# Articulated Body Algorithm，输入为模型、关节角、关节速度、扭矩、重力、步长、步数
import pinocchio as pin
import numpy as np
from os.path import dirname, join, abspath
import yaml
def abaLoop(model, data, q0 = None, v0 = None, tau0 = None, ROBOT = False):
    
    
    # 根据Robot关键字的bool值，选取特定的init.yml文件
    if ROBOT:
        # 这里open的基路径是main.py
        with open(join(dirname(str(abspath(__file__))),'robot_init.yml'), 'r', encoding='utf-8') as f:
            init_config= yaml.load(f.read(), Loader=yaml.FullLoader)
    else:
        with open(join(dirname(str(abspath(__file__))),'init.yml'), 'r', encoding='utf-8') as f:
            init_config= yaml.load(f.read(), Loader=yaml.FullLoader)
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    gravity=init_config['gravity']
    dt=init_config['dt']
    nsteps = init_config['nsteps']


    if q0 is None:
        q0 = pin.neutral(model)
    if v0 is None:
        v0 = np.zeros(model.nv)
    if tau0 is None:
        tau0 = np.zeros(model.nv)
    model.gravity.linear[:2] = 0.0
    model.gravity.linear[2] = gravity
    qs = [q0]
    vs = [v0]
    for i in range(nsteps):
        q = qs[i]
        v = vs[i]
        # 输入模型、关节角、关节速度、扭矩、重力、步长、步数，输出模型的关节加速度
        a1 = pin.aba(model, data, q, v, tau0)
        vnext = v + dt * a1
        qnext = pin.integrate(model, q, dt * vnext)
        qs.append(qnext)
        vs.append(vnext)
    return qs, vs, dt